from data.bbox.bbox_dataset import DetectionDataset
import os
import numpy as np
import json
class retailDetDataset(DetectionDataset):
    def __init__(self, objs):
        super(retailDetDataset,self).__init__()
        self.objs = objs
        self.classes = ["person",]
    def at_with_image_path(self, idx):
        obj = self.objs[idx]
        path = obj[0]
        bboxes = []
        for box in obj[1]:
            xmin = box["minx"]
            ymin = box["miny"]
            xmax = box["maxx"]
            ymax = box["maxy"]
            bboxes.append([xmin,ymin,xmax,ymax,self.classes.index(box["name"],0)])
        return path, np.array(bboxes).astype(float)
    def __len__(self):
        return len(self.objs)

def get_retail_dataset():
    from sklearn.model_selection import train_test_split
    root = "/data1/zyx/yks/dataset/retail/"
    img_root = "/data1/zyx/yks/dataset/retail/train_image/"
    all_imgs = []
    for r,_,names in os.walk(img_root):
        for name in names:
            path = os.path.join(r, name)
            anno_path = path.replace("train_image","train_label").replace(".jpg",".json")
            anno = json.load(open(anno_path,"rb"))["annotation"][0]["object"]
            if len(anno) > 0:
                all_imgs.append((path,anno))
    train_objs,val_objs = train_test_split(all_imgs,random_state=43,test_size=.1)
    train_da = retailDetDataset(train_objs)
    val_da = retailDetDataset(val_objs)
    train_da.viz()
    # train_da.to_coco("/data1/zyx/yks/dataset/retail/annotations/instances_trainretail.json")
    # val_da.to_coco("/data1/zyx/yks/dataset/retail/annotations/instances_valretail.json")
    # train_da.to_roidb("/data1/zyx/yks/dataset/retail/annotations/train.roidb")
    # val_da.to_roidb("/data1/zyx/yks/dataset/retail/annotations/val.roidb")
if __name__ == '__main__':
    get_retail_dataset()